import akshare as ak
import pandas as pd
from datetime import datetime, timedelta
from datetime import datetime, timedelta
import os
from pathlib import Path

class MarketDataService:
    def __init__(self, update_stock_list=False):
        # 缓存股票列表数据，避免频繁请求
        self._stock_list_cache = None
    
        # 设置本地存储路径
        self._data_dir = Path('data/market')
        self._data_dir.mkdir(parents=True, exist_ok=True)
        self._stock_list_file = self._data_dir / 'stock_list.csv'
        
        # 如果需要更新或本地文件不存在，则更新股票列表
        if update_stock_list or not self._stock_list_file.exists():
            self._force_update_stock_list()
        else:
            # 否则尝试从本地文件加载
            self._load_stock_list_from_file()
    
    def _force_update_stock_list(self):
        """强制从API更新股票列表数据"""
        try:
            df = ak.stock_info_a_code_name()
            df.to_csv(self._stock_list_file, index=False)
            print(f"Stock list updated successfully: {len(df)} stocks")
            self._stock_list_cache = df
        except Exception as e:
            print(f"Error updating stock list: {str(e)}")
            self._load_stock_list_from_file()
    
    def _load_stock_list_from_file(self):
        """从本地文件加载股票列表"""
        try:
            if self._stock_list_file.exists():
                df = pd.read_csv(self._stock_list_file, dtype = {"code": str})
                
                print(f"Stock list loaded from file: {len(df)} stocks")
                self._stock_list_cache =  df
            else:
                print("Local stock list file not found")
        except Exception as e:
            print(f"Error loading stock list from file: {str(e)}")
    
    def _get_stock_list(self):
        
        # 如果缓存过期，从本地文件重新加载
        self._load_stock_list_from_file()
        return self._stock_list_cache
    
    def update_stock_list(self):
        """手动触发更新股票列表"""
        self._force_update_stock_list()

    def _get_index_list(self):
        try:
            # 获取指数列表
            df = ak.index_stock_info()
            # 重命名列以匹配股票数据格式
            df = df.rename(columns={
                'index_code': 'code',
                'display_name': 'name'
            })
            # df的code前面加上sh
            df['code'] = 'sh' + df['code']
            # 更新缓存
            self._index_list_cache = df
            return df
        except Exception as e:
            print(f"Error fetching index list: {str(e)}")
            # 如果获取新数据失败但有缓存，返回缓存
            if self._index_list_cache is not None:
                return self._index_list_cache
            raise

    def search_stocks(self, query, limit=30):
        """搜索股票和指数"""
        try:
            results = []
            
            # 搜索股票
            stock_info_df = self._stock_list_cache

             # Ensure we have data to search
            if stock_info_df is None or stock_info_df.empty:
                print("Warning: Stock list is empty")
                return results


            filtered_stocks = stock_info_df[
                (stock_info_df['name'].str.contains(query, case=False)) | 
                (stock_info_df['code'].str.contains(query, case=False))
            ]
            
            # 搜索指数
            index_info_df = self._get_index_list()
            filtered_indices = index_info_df[
                (index_info_df['name'].str.contains(query, case=False)) | 
                (index_info_df['code'].str.contains(query, case=False))
            ]
            
            # 合并结果
            for _, row in filtered_stocks.head(limit//2).iterrows():
                results.append({
                    "symbol": row['code'],
                    "ticker": row['code'],
                    "full_name": row['name'],
                    "description": f"{row['code']} - {row['name']}",
                    "exchange": "SSE" if str(row['code']).startswith('6') else "SZSE",
                    "type": "stock"
                })
                
            for _, row in filtered_indices.head(limit//2).iterrows():
                results.append({
                    "symbol": row['code'],
                    "ticker": row['code'],
                    "full_name": row['name'],
                    "description": row['name'],
                    "exchange": "INDEX",
                    "type": "index"
                })
            
            return results[:limit]
            
        except Exception as e:
            print(f"Error in search: {str(e)}")
            return []

    def get_symbol_info(self, symbol):
        """获取股票配置信息"""
        try:
            # 如果是以sh或sz开头，先尝试查找指数
            if symbol.startswith('sh') or symbol.startswith('sz'):
                index_info_df = self._get_index_list()
                matched_indices = index_info_df[index_info_df['code'] == symbol]
                if not matched_indices.empty:
                    index = matched_indices.iloc[0]
                    return {
                        "name": index['name'],
                        "ticker": symbol,
                        "exchange-traded": "INDEX",
                        "exchange-listed": "INDEX",
                        "timezone": "Asia/Shanghai",
                        "minmov": 1,
                        "minmov2": 0,
                        "pointvalue": 1,
                        "session": "0930-1530",
                        "has_intraday": True,
                        "has_no_volume": False,
                        "description": f"{symbol} - {index['name']}",
                        "type": "index",
                        "supported_resolutions": ["1", "5", "15", "30", "60", "D", "W", "M"],
                        "pricescale": 100,
                        "ticker": symbol,
                        "data_status": "streaming"
                    }
            stock_info_df = self._stock_list_cache
            if stock_info_df is None or stock_info_df.empty:
                print("Warning: Stock list is empty")
                return None
            # 确保代码列是字符串类型
            stock_info_df['code'] = stock_info_df['code'].astype(str)
            
            # 查找匹配的股票
            matched_stocks = stock_info_df[stock_info_df['code'] == str(symbol)]
            if matched_stocks.empty:
                print(f"Symbol not found: {symbol}")
                return None
                
            stock = matched_stocks.iloc[0]
            # 添加正确的前缀
            prefix = 'sh' if str(symbol).startswith('6') else 'sz'
            full_symbol = f"{prefix}{symbol}"
            
            return {
                "name": stock['name'],
                "ticker": symbol,
                "exchange-traded": "SSE" if str(symbol).startswith('6') else "SZSE",
                "exchange-listed": "SSE" if str(symbol).startswith('6') else "SZSE",
                "timezone": "Asia/Shanghai",
                "minmov": 1,
                "minmov2": 0,
                "pointvalue": 1,
                "session": "0930-1530",
                "has_intraday": True,
                "has_no_volume": False,
                "description": f"{symbol} - {stock['name']}",
                "type": "stock",
                "supported_resolutions": ["1", "5", "15", "30", "60", "D", "W", "M"],
                "pricescale": 100,
                "data_status": "streaming"
            }
            
        except Exception as e:
            print(f"Symbol info error: {str(e)}")
            return None

    def get_stock_quote(self, symbol):
        """获取股票或指数实时行情"""
        try:
            if symbol.startswith('sh') or symbol.startswith('sz'):
                # 获取指数实时行情数据
                df = ak.stock_zh_index_spot_sina()
                # 匹配指数代码
                matched_data = df[df['代码'] == symbol]
            else:
                # 获取股票实时行情数据
                df = ak.stock_zh_a_spot_em()
                # 匹配股票代码或名称
                matched_data = df[(df['代码'] == symbol) | (df['名称'] == symbol)]
            
            if matched_data.empty:
                return {
                    "s": "error",
                    "n": symbol,
                    "errmsg": "Symbol not found"
                }
            
            data = matched_data.iloc[0]
            
            return {
                "s": "ok",
                "n": symbol,
                "v": float(data.get('成交量', 0)),
                "t": int(datetime.now().timestamp()),
                "price": float(data['最新价']),
                "open": float(data.get('今开', data['最新价'])),
                "high": float(data.get('最高', data['最新价'])),
                "low": float(data.get('最低', data['最新价'])),
                "close": float(data['最新价']),
                "volume": float(data.get('成交量', 0)),
                "amount": float(data.get('成交额', 0))
            }
        
        except Exception as e:
            print(f"Error fetching stock quote: {str(e)}")
            return {
                "s": "error",
                "n": symbol,
                "errmsg": str(e)
            }
        

    def _get_minute_kline(self, symbol, resolution, from_date, to_date):
        """获取分钟级别K线数据"""
        try:

            if self.is_index(symbol):
                # 如果是指数，使用 ak.stock_zh_index_min_em 获取数据
                df = ak.index_zh_a_hist_min_em(
                    symbol=symbol[2:],
                    period=resolution,
                    start_date=from_date.strftime('%Y%m%d %H:%M:%S'),
                    end_date=to_date.strftime('%Y%m%d %H:%M:%S')
                )
            else:    
                df = ak.stock_zh_a_hist_min_em(
                    symbol=symbol,
                    period=resolution,
                    start_date=from_date.strftime('%Y%m%d %H:%M:%S'),
                    end_date=to_date.strftime('%Y%m%d %H:%M:%S'),
                    adjust="qfq"
                )
            
            if df.empty:
                next_time = int((datetime.now() - timedelta(days=2)).timestamp())
                return {"s": "no_data", "nextTime": next_time}
            
            # 统一处理日期时间格式
            df['时间'] = pd.to_datetime(df['时间'])
            # 分钟级别数据减去8小时
            timestamps = df['时间'].apply(lambda x: int(x.timestamp()) - 8 * 3600).tolist()
            
            return {
                "s": "ok",
                "t": timestamps,
                "o": df['开盘'].astype(float).tolist(),
                "h": df['最高'].astype(float).tolist(),
                "l": df['最低'].astype(float).tolist(),
                "c": df['收盘'].astype(float).tolist(),
                "v": df['成交量'].astype(float).tolist()
            }
            
        except Exception as e:
            print(f"Error getting minute kline: {str(e)}")
            return {"s": "error", "errmsg": str(e)}

    def get_kline_data(self, symbol, resolution, from_time, to_time):
        """获取K线数据"""
        try:
            # 转换时间戳为日期字符串
            from_date = datetime.fromtimestamp(from_time)
            to_date = datetime.fromtimestamp(to_time)
            
            # 分钟级别数据单独处理
            if resolution in ['1', '5', '15', '30', '60']:
                return self._get_minute_kline(symbol, resolution, from_date, to_date)
            
            # 对于日线及以上周期，统一处理
            period = {
                'D': 'daily',
                '1D': 'daily',
                'W': 'weekly',
                '1W': 'weekly',
                'M': 'monthly',
                '1M': 'monthly'
            }.get(resolution, 'daily')
            if self.is_index(symbol):
                # 如果是指数，使用 ak.stock_zh_index_hist_em 获取数据
                df = ak.index_zh_a_hist(
                    symbol=symbol[2:],
                    period=period,
                    start_date=from_date.strftime('%Y%m%d'),
                    end_date=to_date.strftime('%Y%m%d')
                )
            else:

                df = ak.stock_zh_a_hist(
                    symbol=symbol,
                    period=period,
                    start_date=from_date.strftime('%Y%m%d'),
                    end_date=to_date.strftime('%Y%m%d'),
                    adjust="qfq"
                )

            if df.empty:
                return {
                    "s": "ok",
                    "t": [],
                    "o": [],
                    "h": [],
                    "l": [],
                    "c": [],
                    "v": []
                }
            
            # 确保日期列的名称正确
            date_column = '日期'
            df[date_column] = pd.to_datetime(df[date_column])
            timestamps = df[date_column].apply(lambda x: int(x.timestamp())).tolist()
            # 转换数据格式
            return {
                "s": "ok",
                "t": timestamps,
                "o": df['开盘'].astype(float).tolist(),
                "h": df['最高'].astype(float).tolist(),
                "l": df['最低'].astype(float).tolist(),
                "c": df['收盘'].astype(float).tolist(),
                "v": df['成交量'].astype(float).tolist()
            }
            
        except Exception as e:
            print(f"Error in get_kline_data: {str(e)}")
            return {
                "s": "error",
                "errmsg": str(e)
            }
   
    def is_index(self, symbol):
        """判断是否为指数"""
        return symbol.startswith('sh') or symbol.startswith('sz')  